import { CompletionOptions } from "../.."
import { AutocompleteLanguageInfo } from "../constants/AutocompleteLanguageInfo"

// TODO: Do we want to stop completions when reaching a `/src/` string?
const SRC_DIRECTORY = "/src/"
// Starcoder2 tends to output artifacts starting with the letter "t"
const STARCODER2_T_ARTIFACTS = ["t.", "\nt", "<file_sep>"]
const PYTHON_ENCODING = "#- coding: utf-8"
const CODE_BLOCK_END = "```"

// const multilineStops: string[] = [DOUBLE_NEWLINE, WINDOWS_DOUBLE_NEWLINE];
const commonStops = [SRC_DIRECTORY, PYTHON_ENCODING, CODE_BLOCK_END]

export function getStopTokens(
	completionOptions: Partial<CompletionOptions> | undefined,
	_lang: AutocompleteLanguageInfo,
	model: string,
): string[] {
	const stopTokens = [
		...(completionOptions?.stop || []),
		// ...multilineStops,
		...commonStops,
		...(model.toLowerCase().includes("starcoder2") ? STARCODER2_T_ARTIFACTS : []),
		// ...lang.topLevelKeywords.map((word) => `\n${word}`),
	]

	return stopTokens
}
